// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/tool/trace/trace_replay_runner.h>

#include <turbo/flags/flag.h>

#include <pollux/common/file/file_systems.h>
#include <pollux/common/memory/memory.h>
#include <pollux/connectors/hive/hive_connector.h>
#include <pollux/connectors/hive/hive_connector_split.h>
#include <pollux/connectors/hive/hive_data_sink.h>
#include <pollux/connectors/hive/table_handle.h>
#include <pollux/connectors/hive/storage_adapters/abfs/RegisterAbfsFileSystem.h>
#include <pollux/connectors/hive/storage_adapters/gcs/register_gcs_filesystem.h>
#include <pollux/connectors/hive/storage_adapters/hdfs/RegisterHdfsFileSystem.h>
#include <pollux/connectors/hive/storage_adapters/s3fs/RegisterS3FileSystem.h>
#include <pollux/core/plan_node.h>
#include <pollux/dwio/dwrf/register_dwrf_reader.h>
#include <pollux/dwio/dwrf/register_dwrf_writer.h>
#include <pollux/exec/operator_trace_reader.h>
#include <pollux/exec/partition_function.h>
#include <pollux/exec/task_trace_reader.h>
#include <pollux/exec/trace_util.h>
#include <pollux/functions/prestosql/aggregates/register_aggregate_functions.h>
#include <pollux/functions/prestosql/registration/registration_functions.h>
#include <pollux/parse/type_resolver.h>
#include <pollux/serializers/compact_row_serializer.h>
#include <pollux/serializers/unsafe_row_serializer.h>
#include <pollux/tool/trace/aggregation_replayer.h>
#include <pollux/tool/trace/filter_project_replayer.h>
#include <pollux/tool/trace/hash_join_replayer.h>
#include <pollux/tool/trace/operator_replayer_base.h>
#include <pollux/tool/trace/partitioned_output_replayer.h>
#include <pollux/tool/trace/table_scan_replayer.h>
#include <pollux/tool/trace/table_writer_replayer.h>
#include <pollux/type/type.h>

#include <pollux/dwio/parquet/register_parquet_reader.h>
#include <pollux/dwio/parquet/register_parquet_writer.h>


TURBO_FLAG(std::string,
           root_dir,
           "",
           "Root directory where the replayer is reading the traced data, it must be "
           "set");
TURBO_FLAG(bool,
           summary,
           false,
           "Show the summary of the tracing including number of tasks and task ids. "
           "It also print the query metadata including query configs, connectors "
           "properties, and query plan in JSON format.");
TURBO_FLAG(bool, short_summary, false, "Only show number of tasks and task ids");
TURBO_FLAG(std::string, query_id, "", "Specify the target query id which must be set");
TURBO_FLAG(std::string,
           task_id,
           "",
           "Specify the target task id, if empty, show the summary of all the traced "
           "query task.");
TURBO_FLAG(std::string, node_id, "", "Specify the target node id.");
TURBO_FLAG(std::string, driver_ids, "", "A comma-separated list of target driver ids");
TURBO_FLAG(std::string,
           table_writer_output_dir,
           "",
           "Specify output directory of TableWriter.");
TURBO_FLAG(double,
           hive_connector_executor_hw_multiplier,
           2.0,
           "Hardware multipler for hive connector.");
TURBO_FLAG(double,
           driver_cpu_executor_hw_multiplier,
           2.0,
           "Hardware multipler for driver cpu executor.");
TURBO_FLAG(int32_t,
           shuffle_serialization_format,
           0,
           "Specify the shuffle serialization format, 0: presto columnar, 1: compact row, 2: spark unsafe row.");
TURBO_FLAG(std::string,
           memory_arbitrator_type,
           "shared",
           "Specify the memory arbitrator type.");
TURBO_FLAG(uint64_t,
           query_memory_capacity_mb,
           0,
           "Specify the query memory capacity limit in GB. If it is zero, then there is no limit.");
TURBO_FLAG(bool, copy_results, false, "Copy the replaying results.");
TURBO_FLAG(std::string,
           function_prefix,
           "",
           "Prefix for the scalar and aggregate functions.");

namespace kumo::pollux::tool::trace {
    namespace {
        VectorSerde::Kind getVectorSerdeKind() {
            switch (FLAGS_shuffle_serialization_format) {
                case 0:
                    return VectorSerde::Kind::kPresto;
                case 1:
                    return VectorSerde::Kind::kCompactRow;
                case 2:
                    return VectorSerde::Kind::kUnsafeRow;
                default:
                    POLLUX_UNSUPPORTED(
                        "Unsupported shuffle serialization format: {}",
                        static_cast<int>(FLAGS_shuffle_serialization_format));
            }
        }

        void printTaskMetadata(
            const std::string &taskTraceDir,
            memory::MemoryPool *pool,
            std::ostringstream &oss) {
            auto taskMetaReader = std::make_unique<exec::trace::TaskTraceMetadataReader>(
                taskTraceDir, pool);
            const auto queryConfigs = taskMetaReader->queryConfigs();
            const auto connectorProperties = taskMetaReader->connectorProperties();
            const auto queryPlan = taskMetaReader->queryPlan();

            oss << "\n++++++Query configs++++++\n";
            for (const auto &queryConfigEntry: queryConfigs) {
                oss << "\t" << queryConfigEntry.first << ": " << queryConfigEntry.second
                        << "\n";
            }
            oss << "\n++++++Connector configs++++++\n";
            for (const auto &connectorPropertyEntry: connectorProperties) {
                oss << connectorPropertyEntry.first << "\n";
                for (const auto &propertyEntry: connectorPropertyEntry.second) {
                    oss << "\t" << propertyEntry.first << ": " << propertyEntry.second
                            << "\n";
                }
            }
            oss << "\n++++++Task query plan++++++\n";
            oss << queryPlan->toString(true, true);
        }

        void printPipelineTraceSummary(
            const std::string &taskTraceDir,
            const std::string &nodeId,
            uint32_t pipelineId,
            uint32_t driverId,
            memory::MemoryPool *pool,
            std::ostringstream &oss) {
            const auto opTraceDir = exec::trace::getOpTraceDirectory(
                taskTraceDir, nodeId, pipelineId, driverId);
            const auto opTraceSummary =
                    exec::trace::OperatorTraceSummaryReader(
                        exec::trace::getOpTraceDirectory(
                            taskTraceDir, nodeId, pipelineId, driverId),
                        pool)
                    .read();
            oss << "driver " << driverId << ": " << opTraceSummary.toString() << "\n";
        }

        void printTaskTraceSummary(
            const std::string &traceDir,
            const std::string &queryId,
            const std::string &taskId,
            const std::string &nodeId,
            memory::MemoryPool *pool,
            std::ostringstream &oss) {
            auto fs = filesystems::getFileSystem(traceDir, nullptr);
            const auto taskTraceDir =
                    exec::trace::getTaskTraceDirectory(traceDir, queryId, taskId);
            const auto pipelineIds = exec::trace::listPipelineIds(
                exec::trace::getNodeTraceDirectory(taskTraceDir, nodeId), fs);

            oss << "\n++++++Task " << taskId << "++++++\n";
            for (const auto pipelineId: pipelineIds) {
                oss << "\n++++++Pipeline " << pipelineId << "++++++\n";
                const auto driverIds = exec::trace::listDriverIds(
                    exec::trace::getNodeTraceDirectory(taskTraceDir, nodeId),
                    pipelineId,
                    fs);
                for (const auto &driverId: driverIds) {
                    printPipelineTraceSummary(
                        taskTraceDir, nodeId, pipelineId, driverId, pool, oss);
                }
            }
        }

        void printSummary(
            const std::string &rootDir,
            const std::string &queryId,
            const std::string &taskId,
            bool shortSummary,
            memory::MemoryPool *pool) {
            const std::string queryDir =
                    exec::trace::getQueryTraceDirectory(rootDir, queryId);
            const auto fs = filesystems::getFileSystem(queryDir, nullptr);
            const auto taskIds = exec::trace::getTaskIds(rootDir, queryId, fs);
            POLLUX_USER_CHECK(!taskIds.empty(), "No task found under {}", rootDir);

            std::ostringstream summary;
            summary << "\n++++++Query trace summary++++++\n";
            summary << "Number of tasks: " << taskIds.size() << "\n";
            if (shortSummary) {
                summary << "Task ids: " << melon::join("\n", taskIds);
                KLOG(INFO) << summary.str();
                return;
            }

            const auto summaryTaskIds =
                    taskId.empty() ? taskIds : std::vector<std::string>{taskId};
            printTaskMetadata(
                exec::trace::getTaskTraceDirectory(rootDir, queryId, summaryTaskIds[0]),
                pool,
                summary);
            summary << "\n++++++Task Summaries++++++\n";
            for (const auto &taskId: summaryTaskIds) {
                printTaskTraceSummary(
                    rootDir, queryId, taskId, FLAGS_node_id, pool, summary);
            }
            KLOG(INFO) << summary.str();
        }
    } // namespace

    TraceReplayRunner::TraceReplayRunner()
        : cpuExecutor_(std::make_unique<melon::CPUThreadPoolExecutor>(
              std::thread::hardware_concurrency() *
              FLAGS_driver_cpu_executor_hw_multiplier,
              std::make_shared<melon::NamedThreadFactory>(
                  "TraceReplayCpuConnector"))),
          ioExecutor_(std::make_unique<melon::IOThreadPoolExecutor>(
              std::thread::hardware_concurrency() *
              FLAGS_hive_connector_executor_hw_multiplier,
              std::make_shared<melon::NamedThreadFactory>(
                  "TraceReplayIoConnector"))) {
    }

    void TraceReplayRunner::init() {
        POLLUX_USER_CHECK(!FLAGS_root_dir.empty(), "--root_dir must be provided");
        POLLUX_USER_CHECK(!FLAGS_query_id.empty(), "--query_id must be provided");
        POLLUX_USER_CHECK(!FLAGS_node_id.empty(), "--node_id must be provided");

        if (!memory::MemoryManager::testInstance()) {
            memory::MemoryManagerOptions options;
            options.arbitratorKind = FLAGS_memory_arbitrator_type;
            memory::initializeMemoryManager({});
        }
        filesystems::registerLocalFileSystem();
        filesystems::registerS3FileSystem();
        filesystems::registerHdfsFileSystem();
        filesystems::registerGcsFileSystem();
        filesystems::registerAbfsFileSystem();

        dwio::common::registerFileSinks();
        dwrf::registerDwrfReaderFactory();
        dwrf::registerDwrfWriterFactory();

        parquet::registerParquetReaderFactory();
        parquet::registerParquetWriterFactory();

        core::PlanNode::registerSerDe();
        core::ITypedExpr::registerSerDe();
        common::Filter::registerSerDe();
        Type::registerSerDe();
        exec::registerPartitionFunctionSerDe();
        if (!isRegisteredVectorSerde()) {
            serializer::presto::PrestoVectorSerde::registerVectorSerde();
        }
        if (!isRegisteredNamedVectorSerde(VectorSerde::Kind::kPresto)) {
            serializer::presto::PrestoVectorSerde::registerNamedVectorSerde();
        }
        if (!isRegisteredNamedVectorSerde(VectorSerde::Kind::kCompactRow)) {
            serializer::CompactRowVectorSerde::registerNamedVectorSerde();
        }
        if (!isRegisteredNamedVectorSerde(VectorSerde::Kind::kUnsafeRow)) {
            serializer::spark::UnsafeRowVectorSerde::registerNamedVectorSerde();
        }
        connector::hive::HiveTableHandle::registerSerDe();
        connector::hive::LocationHandle::registerSerDe();
        connector::hive::HiveColumnHandle::registerSerDe();
        connector::hive::HiveInsertTableHandle::registerSerDe();
        connector::hive::HiveInsertFileNameGenerator::registerSerDe();
        connector::hive::HiveConnectorSplit::registerSerDe();
        connector::hive::registerHivePartitionFunctionSerDe();
        connector::hive::HiveBucketProperty::registerSerDe();

        functions::prestosql::registerAllScalarFunctions(FLAGS_function_prefix);
        aggregate::prestosql::registerAllAggregateFunctions(FLAGS_function_prefix);
        parse::registerTypeResolver();

        if (!kumo::pollux::connector::hasConnectorFactory("hive")) {
            connector::registerConnectorFactory(
                std::make_shared<connector::hive::HiveConnectorFactory>());
        }

        fs_ = filesystems::getFileSystem(FLAGS_root_dir, nullptr);
        const auto taskTraceDir = exec::trace::getTaskTraceDirectory(
            FLAGS_root_dir, FLAGS_query_id, FLAGS_task_id);
        taskTraceMetadataReader_ =
                std::make_unique<exec::trace::TaskTraceMetadataReader>(
                    taskTraceDir, memory::MemoryManager::getInstance()->tracePool());
    }

    std::unique_ptr<tool::trace::OperatorReplayerBase>
    TraceReplayRunner::createReplayer() const {
        std::unique_ptr<tool::trace::OperatorReplayerBase> replayer;
        const auto taskTraceDir = exec::trace::getTaskTraceDirectory(
            FLAGS_root_dir, FLAGS_query_id, FLAGS_task_id);
        const auto traceNodeName = taskTraceMetadataReader_->nodeName(FLAGS_node_id);
        const auto queryCapacityBytes = (1ULL * FLAGS_query_memory_capacity_mb) << 20;
        if (traceNodeName == "TableWrite") {
            POLLUX_USER_CHECK(
                !FLAGS_table_writer_output_dir.empty(),
                "--table_writer_output_dir is required");
            replayer = std::make_unique<tool::trace::TableWriterReplayer>(
                FLAGS_root_dir,
                FLAGS_query_id,
                FLAGS_task_id,
                FLAGS_node_id,
                traceNodeName,
                FLAGS_driver_ids,
                queryCapacityBytes,
                cpuExecutor_.get(),
                FLAGS_table_writer_output_dir);
        } else if (traceNodeName == "Aggregation") {
            replayer = std::make_unique<tool::trace::AggregationReplayer>(
                FLAGS_root_dir,
                FLAGS_query_id,
                FLAGS_task_id,
                FLAGS_node_id,
                traceNodeName,
                FLAGS_driver_ids,
                queryCapacityBytes,
                cpuExecutor_.get());
        } else if (traceNodeName == "PartitionedOutput") {
            replayer = std::make_unique<tool::trace::PartitionedOutputReplayer>(
                FLAGS_root_dir,
                FLAGS_query_id,
                FLAGS_task_id,
                FLAGS_node_id,
                getVectorSerdeKind(),
                traceNodeName,
                FLAGS_driver_ids,
                queryCapacityBytes,
                cpuExecutor_.get());
        } else if (traceNodeName == "TableScan") {
            const auto connectorId =
                    taskTraceMetadataReader_->connectorId(FLAGS_node_id);
            if (const auto &collectors = connector::getAllConnectors();
                collectors.find(connectorId) == collectors.end()) {
                const auto hiveConnector =
                        connector::getConnectorFactory("hive")->newConnector(
                            connectorId,
                            std::make_shared<config::ConfigBase>(
                                std::unordered_map<std::string, std::string>()),
                            ioExecutor_.get());
                connector::registerConnector(hiveConnector);
            }
            replayer = std::make_unique<tool::trace::TableScanReplayer>(
                FLAGS_root_dir,
                FLAGS_query_id,
                FLAGS_task_id,
                FLAGS_node_id,
                traceNodeName,
                FLAGS_driver_ids,
                queryCapacityBytes,
                cpuExecutor_.get());
        } else if (traceNodeName == "Filter" || traceNodeName == "Project") {
            replayer = std::make_unique<tool::trace::FilterProjectReplayer>(
                FLAGS_root_dir,
                FLAGS_query_id,
                FLAGS_task_id,
                FLAGS_node_id,
                traceNodeName,
                FLAGS_driver_ids,
                queryCapacityBytes,
                cpuExecutor_.get());
        } else if (traceNodeName == "HashJoin") {
            replayer = std::make_unique<tool::trace::HashJoinReplayer>(
                FLAGS_root_dir,
                FLAGS_query_id,
                FLAGS_task_id,
                FLAGS_node_id,
                traceNodeName,
                FLAGS_driver_ids,
                queryCapacityBytes,
                cpuExecutor_.get());
        } else {
            POLLUX_UNSUPPORTED("Unsupported operator type: {}", traceNodeName);
        }
        POLLUX_USER_CHECK_NOT_NULL(replayer);
        return replayer;
    }

    void TraceReplayRunner::run() {
        if (FLAGS_summary || FLAGS_short_summary) {
            auto pool = memory::memoryManager()->addLeafPool("replayer");
            printSummary(
                FLAGS_root_dir,
                FLAGS_query_id,
                FLAGS_task_id,
                FLAGS_short_summary,
                pool.get());
            return;
        }
        POLLUX_USER_CHECK(!FLAGS_task_id.empty(), "--task_id must be provided");
        createReplayer()->run(FLAGS_copy_results);
    }
} // namespace kumo::pollux::tool::trace
